'''
数据处理脚本 seg  windows linux
'''
import datetime
import glob
import json
import os.path
import random
import shutil
import sys

import cv2
import numpy as np

sys.path.append('yolov5')
from yolov5.utils.dataloaders import get_hash # 避免重复转换数据集

class base:
    '''
    labelme格式 to yolo格式
    '''

    def __init__(self):
        pass

    def single_jsontotxt(self, json_path, out_path):
        '''
        seg: json转txt，单个文件，
        :param json_path:输入json文件地址
        :param out_path:输出txt文件地址  cls xywh
        :return:
        '''

        # load json
        with open(json_path, 'r') as load_f:
            content = json.load(load_f)

        # each box
        file_str = ''
        W, H = content['imageWidth'], content['imageHeight']
        for t in content['shapes']:

            # type, (标签名转typeid) 修改它
            # ref = {'yuanzhu_hong': 0, 'yuanzhu_lan': 1, 'fangkuai_hong': 2}
            if self.ref.get(t['label']) is None: # 不在self.ref里的标签不考虑
                continue
            type = self.ref[t['label']]
            file_str += str(type)
            # xys
            for x, y in t['points']:
                x = x / W
                y = y / H
                file_str += ' ' + str(round(x, 6)) + ' ' + str(round(y, 6))
            file_str += '\n'


            # 根据目标随机裁剪，保存裁剪图片和标签 20240124
            if False:
                try:

                    W, H = content['imageWidth'], content['imageHeight']
                    # 1 获取裁剪框
                    # xys = t['points'] # [(x,y)...]
                    xys = np.array(t['points']) # n*2
                    x1, y1, x2, y2 = min(xys[:,0]), min(xys[:,1]), max(xys[:,0]), max(xys[:,1]),
                    bw,bh = x2-x1, y2-y1

                    scale = 1.0
                    len_max = max(bw, bh)
                    x1_rand = random.randint(0, int(scale*len_max))
                    x2_rand = random.randint(0, int(scale * len_max))
                    y1_rand = random.randint(0, int(scale * len_max))
                    y2_rand = random.randint(0, int(scale * len_max))

                    x1_clip = int(max(x1 - x1_rand, 0))
                    x2_clip = int(min(x2 + x2_rand, W))
                    y1_clip = int(max(y1 - y1_rand, 0))
                    y2_clip = int(min(y2 + y2_rand, H))

                    # 2 裁剪图片和标签
                    img = cv2.imread(json_path[:-5]+'.jpg')
                    img_clip = img[y1_clip:y2_clip,x1_clip:x2_clip] # 图片
                    type = self.ref[t['label']]
                    lab_str = str(type) # 标签， 不考虑目标标签外的标签
                    for x, y in t['points']:
                        x = (x - x1_clip) / (x2_clip-x1_clip)
                        y = (y - y1_clip) / (y2_clip-y1_clip)
                        lab_str += ' ' + str(round(x, 6)) + ' ' + str(round(y, 6))

                    # 3 保存
                    js_name = json_path.split('\\')[-1] # wind
                    save_dir = json_path[:-len(js_name)-2]+'_mini' # window
                    if not os.path.exists(save_dir):
                        os.mkdir(save_dir)
                    cv2.imwrite(f'{save_dir}/{js_name[:-5]}_clip.jpg', img_clip)
                    fp = open(f'{save_dir}/{js_name[:-5]}_clip.txt', mode="w", encoding="utf-8")  # win
                    fp.write(lab_str)
                    fp.close()
                except Exception as e:
                    print(e)
                    print(json_path)
            # # cat
            # file_str += str(type) + ' ' + str(round(x, 6)) + ' ' + str(round(y, 6)) + ' ' + str(
            #     round(w, 6)) + ' ' + str(round(h, 6)) + '\n'  # 4月12日更改
        # if len(file_str)==0:
        #     return

        # save
        filename = out_path
        if os.path.exists(filename):
            os.remove(filename)
        try:
            os.mknod(filename)  # win不支持
            fp = open(filename, mode="r+", encoding="utf-8")
        except:
            print('非 linux平台')
            fp = open(filename, mode="w", encoding="utf-8") # win
        fp.write(file_str[:-1])
        fp.close()

    def do_augment(self, json_path, isClip=True, scale = 1.0):
        # load json
        with open(json_path, 'r') as load_f:
            content = json.load(load_f)
        W, H = content['imageWidth'], content['imageHeight']

        if isClip:
            try:
                for ind, t in enumerate(content['shapes']): # each target
                    if self.ref.get(t['label']) is None:  # 不在self.ref里的标签不考虑
                        continue
                    # 1 获取裁剪框
                    # xys = t['points'] # [(x,y)...]
                    xys = np.array(t['points'])  # n*2
                    x1, y1, x2, y2 = min(xys[:, 0]), min(xys[:, 1]), max(xys[:, 0]), max(xys[:, 1]),
                    bw, bh = x2 - x1, y2 - y1

                    # scale = 1.0
                    len_max = max(bw, bh)
                    if random.random() > 1:# 策略1
                        x1_rand = random.randint(0, int(scale * len_max))
                        x2_rand = random.randint(0, int(scale * len_max))
                        y1_rand = random.randint(0, int(scale * len_max))
                        y2_rand = random.randint(0, int(scale * len_max))
                    else:
                        # 策略2 20240510
                        x1_rand = len_max/8
                        x2_rand = len_max/8
                        y1_rand = len_max/8
                        y2_rand = len_max/8
                    if t['label'] == 'bar': # 策略3 20240511
                        continue

                    x1_clip = int(max(x1 - x1_rand, 0))
                    x2_clip = int(min(x2 + x2_rand, W))
                    y1_clip = int(max(y1 - y1_rand, 0))
                    y2_clip = int(min(y2 + y2_rand, H))

                    # 2 裁剪图片和标签
                    img = cv2.imread(json_path[:-5] + '.jpg')
                    img_clip = img[y1_clip:y2_clip, x1_clip:x2_clip]  # 保存裁剪图片

                    lab_str = ''
                    for st in content['shapes']: # 统计裁切框内的所有标签
                        # type = self.ref[t['label']]
                        if self.ref.get(st['label']) is None:  # 不在self.ref里的标签不考虑
                            continue
                        type = self.ref[st['label']]  # bug修正 20240305
                        # 判断目标是否在img_clip外
                        cnt = 0
                        for x, y in st['points']:
                            x = (x - x1_clip) / (x2_clip - x1_clip)
                            y = (y - y1_clip) / (y2_clip - y1_clip)
                            if x < 0 or x > 1 or y < 0 or y > 1:  # 20250125 在clip图片外
                                cnt += 1
                        # if cnt > len(st['points'])*0.5 : #如果目标有一半的点在clip图片外，不考虑该目标
                        #     continue
                        if cnt > len(st['points'])*0.0 : #如果目标有点在clip图片外，不考虑该目标
                            continue
                        lab_str += str(type)  # 记录标签，
                        for x, y in st['points']:
                            x = (x - x1_clip) / (x2_clip - x1_clip)
                            y = (y - y1_clip) / (y2_clip - y1_clip)
                            x = min(max(x,0),1)
                            y = min(max(y,0),1) # 0-1

                            lab_str += ' ' + str(round(x, 6)) + ' ' + str(round(y, 6))
                        lab_str += '\n'
                    # 3 保存
                    # js_name = json_path.split('\\')[-1]  # wind
                    # save_dir = json_path[:-len(js_name) - 1] + '_mini'  # window
                    # wind or linux
                    js_name = os.path.basename(json_path)
                    save_dir = os.path.dirname(json_path) + '_mini'

                    if not os.path.exists(save_dir):
                        os.mkdir(save_dir)
                    cv2.imwrite(f'{save_dir}/{js_name[:-5]}_clip_{ind}.jpg', img_clip)
                    fp = open(f'{save_dir}/{js_name[:-5]}_clip_{ind}.txt', mode="w", encoding="utf-8")  # win
                    fp.write(lab_str[:-1])
                    fp.close()
            except Exception as e:
                print(e)
                print(json_path)

    def make_txt_from_json(self):
        '''
        json to txt
        '''
        print(f'Make_txt_from_json...')
        ls = glob.glob(self.json)
        total = 0
        for i in ls:
            txt_path = f'{i[:-5]}.txt'
            self.single_jsontotxt(i, txt_path)  # 转换
            # self.do_augment(i) # 随机裁剪
            total += 1
            # print(f'{total}/{int(len(ls))}: saved to {txt_path}')
        print(f'Make_txt_from_json , Done.')

    def augment_from_json(self, random_scale = 1.0): # 随机裁剪
        '''
        随机裁剪,生成 ~_mini文件夹
        '''
        print(f'Augment: random clip...')
        ls = glob.glob(self.json)
        total = 0
        for i in ls:
            self.do_augment(i, scale = random_scale) # 随机裁剪
            total += 1
            # print(f'{total}/{int(len(ls))}: saved to {txt_path}')
        print(f'Augment: random clip, Done.')
    # def make_yolov5(self):
    #     pass

    def update_yolov5(self):
        '''
        make yolov5 or add a unit to yolov5
        '''

        # source
        self.info()

        # dst
        if not os.path.exists(self.yolov5):
            os.makedirs(self.yolov5)
        else:
            print(f'rm {self.yolov5}/*/*...')
            if os.path.exists(f'{self.yolov5}images'):
                shutil.rmtree(f'{self.yolov5}images') # 删除，重新转换,目前不能增量式添加数据
            if os.path.exists(f'{self.yolov5}labels'):
                shutil.rmtree(f'{self.yolov5}labels')
        out_fir_tra_img = self.yolov5 + 'images/train/'
        out_fir_val_img = self.yolov5 + 'images/val/'
        out_fir_tra_txt = self.yolov5 + 'labels/train/'
        out_fir_val_txt = self.yolov5 + 'labels/val/'
        iltv_dir = [out_fir_tra_img, out_fir_val_img, out_fir_tra_txt, out_fir_val_txt]
        for i in iltv_dir:
            if not os.path.exists(i): # 如果不存在
                os.makedirs(i)

        print('before:')
        t_i, v_i, t_tx, v_tx = len(os.listdir(out_fir_tra_img)), len(os.listdir(out_fir_val_img)), \
                               len(os.listdir(out_fir_tra_txt)), len(os.listdir(out_fir_val_txt))
        print(f'{out_fir_tra_img}:{t_i}')
        print(f'{out_fir_val_img}:{v_i}')
        print(f'{out_fir_tra_txt}:{t_tx}')
        print(f'{out_fir_val_txt}:{v_tx}')

        # 划分 各类随机抽1%
        ls_img = glob.glob(self.jpg)
        random.shuffle(ls_img)
        flag = max(int(0.2 * len(ls_img)) + 1, 1)  # 20%
        # ls_img_val = ls_img[:flag]
        # ls_img_tra = ls_img[flag:]
        ls_img_val = ls_img[:flag]
        ls_img_tra = ls_img[0:]

        # 添加至验证集

        for i in ls_img_val:
            # jpg
            # nam = i.split('/')[-1]
            # nam = i.split('\\')[-1] # windows
            # self.buff = i.split('\\')[-2] # user
            nam = os.path.basename(i) # wind, linux
            self.buff = os.path.basename(os.path.dirname(i)) # wind, linux
            shutil.copy(i, out_fir_val_img + f'{self.buff}_{nam}')  # 添加至val_img
            # lab
            # txt_abpath = self.txt + nam[:-4] + '.txt'  #
            txt_abpath = f'{i[:-4]}.txt'  #
            if os.path.exists(txt_abpath):  # 如果有标签文件
                shutil.copy(txt_abpath, out_fir_val_txt + f'{self.buff}_{nam[:-4]}.txt')

        # 添加至训练集
        for ind, i in enumerate(ls_img_tra):
            # jpg
            # nam = i.split('/')[-1]
            # nam = i.split('\\')[-1]  # windows
            # self.buff = i.split('\\')[-2]  # user 父文件名为前缀
            nam = os.path.basename(i) # wind, linux
            self.buff = os.path.basename(os.path.dirname(i)) # wind, linux
            shutil.copy(i, out_fir_tra_img + f'{self.buff}_{nam}')  # 添加至tra_img
            # lab
            txt_abpath = f'{i[:-4]}.txt'
            if os.path.exists(txt_abpath):
                shutil.copy(txt_abpath, out_fir_tra_txt + f'{self.buff}_{nam[:-4]}.txt')  # 添加至tra_lab

        # check
        print('after:')
        t_i_d, v_i_d, t_tx_d, v_tx_d = len(os.listdir(out_fir_tra_img)), len(os.listdir(out_fir_val_img)), \
                                       len(os.listdir(out_fir_tra_txt)), len(os.listdir(out_fir_val_txt))
        print(f'{out_fir_tra_img}:{t_i_d}')
        print(f'{out_fir_val_img}:{v_i_d}')
        print(f'img added: {t_i_d + v_i_d - t_i - v_i}')
        print(f'{out_fir_tra_txt}:{t_tx_d}')
        print(f'{out_fir_val_txt}:{v_tx_d}')
        print(f'lab added: {t_tx_d + v_tx_d - t_tx - v_tx}')

        print(f'Update_yolov5, Done.')

    def remove_yolov5(self, buff):
        '''
        remove a unit from yolov5
        :return:
        '''
        ls = glob.glob(f'{self.yolov5}*/*/*')
        for ind, i in enumerate(ls):
            # if i.split('/')[-1][:len(buff)] == buff:  # 根据前缀删除
            name = os.path.basename(i) # wind, linux
            if name[:len(buff)] == buff:  # windows
                os.remove(i)
                print(f'{ind}/{len(ls)}:{i}')

    def info(self):
        # ls = glob.glob(self.glob_str)
        # print(self.buff)
        # print(
        #     f'origin:{len(ls)}\njpg:{len(os.listdir(self.jpg))}\njson:{len(os.listdir(self.json))}\ntxt:{len(os.listdir(self.txt))}\n')
        pass

    def info_txt(self):
        pass

    def resize_imgs(self, target_size = 416): # 缩小训练集图片分辨率，加快训练

        imgs_glob = self.yolov5+'images/*/*.jpg'
        ls = glob.glob(imgs_glob)
        for i in ls:
            img = cv2.imread(i)
            scale = target_size/max(img.shape[:2])
            # print(scale)
            try:
                img = cv2.resize(img, dsize=None,fx=scale,fy=scale)
            except:
                print(i)
            cv2.imwrite(i, img)
        print(f'resized done, total {len(ls)}, scale {scale}')



# class labelme2yoloseg(base):
#     jpg = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.jpg'
#     json = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.json'
#     txt = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.txt'
#     yolov5 = r'D:\data\231207huoni\trainseg_zhihe\imglabs_yolov5/'  # yolov5数据集地址
#     buff = 'd2024_0102'
#     # ref = {'zhihe': 0, 'zhuji_zheng': 1, 'zhuji_fan': 2, 'dianchi_zheng': 3, 'dianchi_fan': 4, 'xianshu': 5, 'shuomingshu': 6, 'wandai': 7 }
#     ref = {'zhihe': 0, 'logo': 1 }
#
#
#     def run(self):
#         self.make_txt_from_json()
#         self.update_yolov5()
#         # self.info()
#         self.resize_imgs(scale=0.3) # (2048, 3072, 3)


class labelme2yoloseg(base):
    def __init__(self):
        self.jpg = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.jpg'
        self.json = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.json'
        self.txt = r'D:\data\231207huoni\trainseg_zhihe\add_imgs\NOREAD_20231230_150531_407\*.txt'
        self.yolov5 = r'D:\data\231207huoni\trainseg_zhihe\imglabs_yolov5/'  # yolov5数据集地址
        self.buff = 'd2024_0102'
        # ref = {'zhihe': 0, 'zhuji_zheng': 1, 'zhuji_fan': 2, 'dianchi_zheng': 3, 'dianchi_fan': 4, 'xianshu': 5, 'shuomingshu': 6, 'wandai': 7 }
        self.ref = {'zhihe': 0, 'logo': 1}
        self.scale = 1.0

    def run(self, isAug = False, random_scale = 1.0):
        self.make_txt_from_json()
        if isAug:
            self.augment_from_json(random_scale = random_scale) # 随机裁剪
        self.update_yolov5()
        # self.info()
        self.resize_imgs(target_size=self.target_size)  # (2048, 3072, 3)

if __name__ == '__main__':
    # deepcam_baby().remove_yolov5(buff='d03_03')

    # video2imgs().video2imgs_muti()
    # video2imgs().get_half()

    # labelme2yoloseg().run()

    root_path = r'D:\data\231207huoni'
    name = 'strap'
    # 1 转换格式 labelme转为yolo数据集
    l2y = labelme2yoloseg()
    l2y.jpg = fr'{root_path}\trainseg_{name}\add_imgs\*\*.jpg'
    l2y.json = fr'{root_path}\trainseg_{name}\add_imgs\*\*.json'
    l2y.txt = fr'{root_path}\trainseg_{name}\add_imgs\*\*.txt'
    l2y.cls = fr'{root_path}\trainseg_{name}\add_imgs\classes.txt'
    l2y.md5 = fr'{root_path}\trainseg_{name}\add_imgs\md5.txt'  # jpg和json文件
    l2y.target_size = 640  # (2048, 3072, 3)

    l2y.yolov5 = fr'{root_path}\trainseg_{name}\imglabs_yolov5/'  # yolov5数据集地址
    l2y.buff = datetime.datetime.now().strftime('d%Y%m%d_')  # 前缀

    # 读取label
    l2y.ref = {}
    # l2y.ref = {'zhihe': 0, 'logo': 1}  #
    with open(l2y.cls) as f:
        lbs = f.read().strip().splitlines()
    for ind, lb in enumerate(lbs):
        l2y.ref[lb.strip()] = ind

    # 转换
    isUpdate = False  # 是否转换数据集
    jpg_paths = glob.glob(fr'{root_path}\trainseg_{name}\add_imgs\*[!mini]\*.jpg')  # 不包含mini数据
    json_paths = glob.glob(fr'{root_path}\trainseg_{name}\add_imgs\*[!mini]\*.json')
    cur_md5_value = get_hash(jpg_paths + json_paths)
    if not os.path.exists(l2y.md5):
        with open(l2y.md5, 'w') as fp:
            fp.write('')
        # os.mknod(l2y.md5) # wind
    with open(l2y.md5, 'r') as fp:
        pre_md5_Value = fp.read()
    if not pre_md5_Value == cur_md5_value:
        with open(l2y.md5, 'w') as fp:
            fp.write(cur_md5_value)
        isUpdate = True
    print(f'原始图片和json文件是否变化（不含_mini） {isUpdate}')
    if isUpdate:
        l2y.run()  # 转换